# [Language Agent Tree Search](https://arxiv.org/abs/2310.04406) (LATS), by Zhou, et. al, is a general LLM agent search algorithm that combines reflection/evaluation and search (specifically monte-carlo trees search) to get achieve better overall task performance compared to similar techniques like ReACT, Reflexion, or Tree of Thoughts.

from ApiTools import apiBase,apiTools
import math
from collections import deque
from typing import Optional
from langchain_core.output_parsers.openai_tools import (
    JsonOutputToolsParser,
    PydanticToolsParser,
)

from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage,SystemMessage
from pydantic import BaseModel, Field
from dotenv import load_dotenv
#load_dotenv(sys.argv[1],override=True)
load_dotenv()

sys_prompt = apiBase.argv_json(1,[{"role": "system", "content": "你是聊天助手"},{"role": "user", "content": "2+2"},{"role": "assistant", "content": "4"}])
#print(f'sys_prompt={sys_prompt}')
question = apiBase.argv(2,"1+1等于多少?")
#sys_prompt.append({"role": "user", "content": question})
#print(f'question={question}')
toolInfo = apiBase.argv_json(4,{"api_name":"searchtool","remark":"search tool","api_url":"api","type":"print"})
tavily_tool = apiTools.create_tool(toolInfo["api_name"],toolInfo["remark"],toolInfo["api_url"],toolInfo["type"])
llm = apiTools.llm

lsTmpl=[]
for obj in sys_prompt:
    if "system" == obj["role"]:
        lsTmpl.append( ('system',obj["content"] ) )
    if "user" == obj["role"]:
        lsTmpl.append( ('human',obj["content"] ) )
    if "human" == obj["role"]:
        lsTmpl.append( ('human',obj["content"] ) )
    if "assistant" == obj["role"]:
        lsTmpl.append( ('ai',obj["content"] ) )
    if "ai" == obj["role"]:
        lsTmpl.append( ('ai',obj["content"] ) )
lsTmpl.append(("user", "{input}"))
lsTmpl.append(MessagesPlaceholder(variable_name="messages", optional=True))
prompt_template = ChatPromptTemplate.from_messages(lsTmpl)
        
initial_answer_chain = prompt_template | llm
#res = initial_answer_chain.invoke({"input": question})
parser = JsonOutputToolsParser(return_id=True)

# ## Define Language Agent
# 
# Our agent will have three primary LLM-powered processes:
# 1. Reflect: score the action based on the tool response.
# 2. Initial response: to create the root node and start the search.
# 3. Expand: generate 5 candidate "next steps" from the best spot in the current tree
# 
# For more "Grounded" tool applications (such as code synthesis), you could integrate code execution into the reflection/reward step. This type of external feedback is very useful (though adds complexity to an already complicated example notebook).

# %%

from langgraph.prebuilt import ToolNode


tools = [tavily_tool]
tool_node = ToolNode(tools=tools)

class Reflection(BaseModel):
    reflections: str = Field(
        description="The critique and reflections on the sufficiency, superfluency,"
        " and general quality of the response"
    )
    score: int = Field(
        description="Score from 0-10 on the quality of the candidate response.",
        gte=0,
        lte=10,
    )
    found_solution: bool = Field(
        description="Whether the response has fully solved the question or task."
    )
    def as_message(self):
        return HumanMessage(
            content=f"Reasoning: {self.reflections}\nScore: {self.score}"
        )

    @property
    def normalized_score(self) -> float:
        return self.score / 10.0

class Node:
    def __init__(
        self,
        messages: list[BaseMessage],
        reflection: Reflection,
        parent: Optional["Node"] = None,
    ):
        self.messages = messages
        self.parent = parent
        self.children = []
        self.value = 0
        self.visits = 0
        self.reflection = reflection
        self.depth = parent.depth + 1 if parent is not None else 1
        self._is_solved = reflection.found_solution if reflection else False
        if self._is_solved:
            self._mark_tree_as_solved()
        self.backpropagate(reflection.normalized_score)

    def __repr__(self) -> str:
        return (
            f"<Node value={self.value}, visits={self.visits},"
            f" solution={self.messages} reflection={self.reflection}/>"
        )

    @property
    def is_solved(self):
        """If any solutions exist, we can end the search."""
        return self._is_solved

    @property
    def is_terminal(self):
        return not self.children

    @property
    def best_child_score(self):
        """Return the child with the highest value."""
        if not self.children:
            return None
        return max(self.children, key=lambda child: int(child.is_solved) * child.value)

    @property
    def height(self) -> int:
        """Check for how far we've rolled out the tree."""
        if self.children:
            return 1 + max([child.height for child in self.children])
        return 1

    def upper_confidence_bound(self, exploration_weight=1.0):
        """Return the UCT score. This helps balance exploration vs. exploitation of a branch."""
        if self.parent is None:
            raise ValueError("Cannot obtain UCT from root node")
        if self.visits == 0:
            return self.value
        # Encourages exploitation of high-value trajectories
        average_reward = self.value / self.visits
        # Encourages exploration of less-visited trajectories
        exploration_term = math.sqrt(math.log(self.parent.visits) / self.visits)
        return average_reward + exploration_weight * exploration_term

    def backpropagate(self, reward: float):
        """Update the score of this node and its parents."""
        node = self
        while node:
            node.visits += 1
            node.value = (node.value * (node.visits - 1) + reward) / node.visits
            node = node.parent

    def get_messages(self, include_reflections: bool = True):
        if include_reflections:
            return self.messages + [self.reflection.as_message()]
        return self.messages

    def get_trajectory(self, include_reflections: bool = True) -> list[BaseMessage]:
        """Get messages representing this search branch."""
        messages = []
        node = self
        while node:
            messages.extend(
                node.get_messages(include_reflections=include_reflections)[::-1]
            )
            node = node.parent
        # Reverse the final back-tracked trajectory to return in the correct order
        return messages[::-1]  # root solution, reflection, child 1, ...

    def _get_all_children(self):
        all_nodes = []
        nodes = deque()
        nodes.append(self)
        while nodes:
            node = nodes.popleft()
            all_nodes.extend(node.children)
            for n in node.children:
                nodes.append(n)
        return all_nodes

    def get_best_solution(self):
        """Return the best solution from within the current sub-tree."""
        all_nodes = [self] + self._get_all_children()
        best_node = max(
            all_nodes,
            # We filter out all non-terminal, non-solution trajectories
            key=lambda node: int(node.is_terminal and node.is_solved) * node.value,
        )
        return best_node

    def _mark_tree_as_solved(self):
        parent = self.parent
        while parent:
            parent._is_solved = True
            parent = parent.parent

from typing_extensions import TypedDict


class TreeState(TypedDict):
    # The full tree
    root: Node
    # The original input
    input: str

# ### Reflection
# 
# The reflection chain will score agent outputs based on the decision and the tool responses.
# We will call this within the other two nodes.

# %%

from langchain_core.runnables import chain as as_runnable

prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "Reflect and grade the assistant response to the user question below.",
        ),
        ("user", "{input}"),
        MessagesPlaceholder(variable_name="candidate"),
    ]
)
reflection_llm_chain = (
    prompt
    | llm.bind_tools(tools=[Reflection], tool_choice="Reflection").with_config(
        run_name="Reflection"
    )
    | PydanticToolsParser(tools=[Reflection])
)

@as_runnable
def reflection_chain(inputs) -> Reflection:
    tool_choices = reflection_llm_chain.invoke(inputs)
    # 处理长度为0的情况
    if(len(tool_choices) == 0):
        reflection = Reflection(reflections="genera",score=6,found_solution=True)
    else:
        reflection = tool_choices[0]
    if not isinstance(inputs["candidate"][-1], AIMessage):
        reflection.found_solution = False
    return reflection

# ### Initial Response
# 
# We start with a single root node, generated by this first step. It responds to the user input either with a tool invocation or a response.

# %%
from langchain_core.prompt_values import ChatPromptValue
from langchain_core.runnables import RunnableConfig

# %%
# Define the node we will add to the graph
def generate_initial_response(state: TreeState) -> dict:
    """Generate the initial candidate response."""
    res = initial_answer_chain.invoke({"input": state["input"]})
    parsed = parser.invoke(res)
    # 处理长度为0的情况
    if len(parsed) > 0:    
        tool_responses = [tool_node.invoke(
            {"messages": [
                AIMessage(content="",tool_calls=[{"name":r["type"], "args":r["args"], 'id':r['id']}]) 
            ]}
        ) for r in parsed]
        output_messages = [res] + [tr['messages'][0] for tr in tool_responses]
        reflection = reflection_chain.invoke(
            {"input": state["input"], "candidate": output_messages}
        )
    else:
        output_messages = [res]
        reflection = reflection_chain.invoke(
            {"input": state["input"], "candidate": output_messages}
        )
        reflection.found_solution=True
    root = Node(output_messages, reflection=reflection)
    return {
        **state,
        "root": root,
    }

def generate_candidates(messages: ChatPromptValue, config: RunnableConfig):
    n = config["configurable"].get("N", 5)
    bound_kwargs = llm.bind_tools(tools=tools).kwargs
    chat_result = llm.generate(
        [messages.to_messages()],
        n=n,
        callbacks=config["callbacks"],
        run_name="GenerateCandidates",
        **bound_kwargs,
    )
    return [gen.message for gen in chat_result.generations[0]]

# 把问题展开回答
expansion_chain = prompt_template | generate_candidates
# res = expansion_chain.invoke({"input": "Write a research report on lithium pollution."})
# res

# #### Candidate generation node
# 
# We will package the candidate generation and reflection steps in the following "expand" node.
# We do all the operations as a batch process to speed up execution.

# %%
from collections import defaultdict

def select(root: Node) -> dict:
    """Starting from the root node a child node is selected at each tree level until a leaf node is reached."""

    if not root.children:
        return root
    
    node = root
    while node.children:
        max_child = max(node.children, key=lambda child: child.upper_confidence_bound())
        node = max_child

    return node

def expand(state: TreeState, config: RunnableConfig) -> dict:
    """Starting from the "best" node in the tree, generate N candidates for the next step."""
    root = state["root"]
    best_candidate: Node = select(root)
    messages = best_candidate.get_trajectory()
    # Generate N candidates from the single child candidate
    new_candidates = expansion_chain.invoke(
        {"input": state["input"], "messages": messages}, config
    )
    parsed = parser.batch(new_candidates)
    flattened = [
        (i, tool_call)
        for i, tool_calls in enumerate(parsed)
        for tool_call in tool_calls
    ]
    tool_responses = [(i,tool_node.invoke(
        {"messages":
            [AIMessage(content="",tool_calls=[{"name":tool_call["type"], "args":tool_call["args"], 'id':tool_call['id']}])]
        }
    )) for i, tool_call in flattened]
    collected_responses = defaultdict(list)
    for i, resp in tool_responses:
        collected_responses[i].append(
            resp['messages'][0]
        )
    output_messages = []
    for i, candidate in enumerate(new_candidates):
        output_messages.append([candidate] + collected_responses[i])

    # Reflect on each candidate
    # For tasks with external validation, you'd add that here.
    reflections = reflection_chain.batch(
        [{"input": state["input"], "candidate": msges} for msges in output_messages],
        config,
    )
    # Grow tree
    child_nodes = [
        Node(cand, parent=best_candidate, reflection=reflection)
        for cand, reflection in zip(output_messages, reflections)
    ]
    best_candidate.children.extend(child_nodes)
    # We have already extended the tree directly, so we just return the state
    return state

# ## Create Graph
# 
# With those two nodes defined, we are ready to define the graph. After each agent step, we have the option of finishing.

# %%
from typing import Literal
from langgraph.graph import END, StateGraph, START
def should_loop(state: TreeState) -> Literal["expand", "__end__"]:
    """Determine whether to continue the tree search."""
    root = state["root"]
    if root.is_solved:
        return END
    if root.height > 5:
        return END
    return "expand"

builder = StateGraph(TreeState)
builder.add_node("start", generate_initial_response)
builder.add_node("expand", expand)
builder.add_edge(START, "start")
builder.add_conditional_edges(
    "start",
    # Either expand/rollout or finish
    should_loop,
)
builder.add_conditional_edges(
    "expand",
    # Either continue to rollout or finish
    should_loop,
)

graph = builder.compile()
# %%
#question = "Generate a table with the average size and weight, as well as the oldest recorded instance for each of the top 5 most common birds."
last_step = None
for step in graph.stream({"input": question}):
    last_step = step
    step_name, step_state = next(iter(step.items()))
    # print(step_name)
    # print("rolled out: ", step_state["root"].height)
    # print("---")

# %%
if "expand" in last_step:
    solution_node = last_step["expand"]["root"].get_best_solution()
else:
    solution_node = last_step["start"]["root"].get_best_solution()
best_trajectory = solution_node.get_trajectory(include_reflections=False)
print(best_trajectory[-1].content)

